#!/bin/bash
# 自动GPU vLLM推理启动脚本：先启动vLLM服务，然后启动Gradio Web界面

#开启V1
export VLLM_USE_V1=1
# 增加vLLM内部超时设置
export VLLM_WORKER_STARTUP_TIMEOUT=600
export VLLM_ENGINE_STARTUP_TIMEOUT=600
export ZMQ_POLL_TIMEOUT=600000

# vLLM服务配置
MODEL_PATH="/gemini/pretrain/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct"  # 模型路径 - 根据实际情况修改
PORT_VLLM=8000                 # vLLM服务端口
LOG_DIR="./logs"               # 日志目录
MAX_TOKENS=2048                # 最大生成token数
API_KEY="ks-proj-1234567890"   # API密钥(可选，如不需要可留空)

# Web界面配置
PORT_WEB=7860                  # Web界面端口
ADAPTER_PATH="./mistral-sft-lora-deepspeed"  # LoRA适配器路径

# 创建日志目录
mkdir -p $LOG_DIR

# ======== 清理环境 ========
echo "停止所有相关进程..."
pkill -f "python.*vllm" || true
pkill -f "api_server" || true
sleep 2

# ======== 环境变量配置 ========
# 自动使用所有可用GPU
unset CUDA_VISIBLE_DEVICES
# 清除其他可能干扰的环境变量
unset NCCL_DEBUG
unset TORCH_DISTRIBUTED_DEBUG
unset TORCH_CPP_LOG_LEVEL

# ======== 检查GPU ========
echo "检测GPU信息..."
nvidia-smi --query-gpu=index,name,memory.total,memory.free --format=csv,noheader

# 计算可用GPU数量
GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)
echo "检测到 $GPU_COUNT 个可用GPU"

# 如果没有GPU，退出
if [ $GPU_COUNT -eq 0 ]; then
    echo "错误：未检测到可用的GPU！"
    exit 1
fi

# ======== 检查vLLM版本 ========
echo "检查vLLM版本..."
python -c "import vllm; print(f'vLLM版本: {vllm.__version__}')"

# ======== 启动vLLM服务 ========
echo "正在启动vLLM服务，使用模型: $MODEL_PATH (检测到 $GPU_COUNT 个GPU)"

# 直接使用检测到的所有GPU
TENSOR_PARALLEL_SIZE=$GPU_COUNT
echo "使用GPU并行度: $TENSOR_PARALLEL_SIZE (全部可用GPU)"

# 使用vLLM 0.8.5支持的参数
python -m vllm.entrypoints.openai.api_server \
  --model $MODEL_PATH \
  --host 0.0.0.0 \
  --port $PORT_VLLM \
  --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
  --max-model-len $MAX_TOKENS \
  --trust-remote-code \
  --api-key "$API_KEY" \
  --seed 42 \
  --gpu-memory-utilization 0.95 \
  --max-num-batched-tokens 16384 \
  --max-num-seqs 128 \
  --disable-log-stats \
  --dtype bfloat16 \
  > $LOG_DIR/vllm_$(date +%Y%m%d).log 2>&1 &

# 保存进程ID
VLLM_PID=$!
echo $VLLM_PID > $LOG_DIR/vllm.pid
echo "vLLM服务已启动，进程ID: $VLLM_PID"
echo "日志保存在: $LOG_DIR/vllm_$(date +%Y%m%d).log"

# ======== 智能等待服务启动 ========
echo "等待vLLM服务启动..."

# 定义函数：检查日志中是否有加载进度信息
check_loading_progress() {
  if grep -q "Loading safetensors checkpoint shards:" $LOG_DIR/vllm_$(date +%Y%m%d).log; then
    local progress=$(grep "Loading safetensors checkpoint shards:" $LOG_DIR/vllm_$(date +%Y%m%d).log | tail -1)
    echo "模型加载进度: $progress"
    return 0
  fi
  return 1
}

# 定义函数：检查日志中是否有成功启动的信息
check_startup_complete() {
  if grep -q "Application startup complete" $LOG_DIR/vllm_$(date +%Y%m%d).log; then
    return 0
  fi
  return 1
}

# 等待服务启动，最多等待15分钟
max_wait_time=900  # 15分钟
start_time=$(date +%s)
echo "开始等待服务启动，最长等待时间: 15分钟..."

while true; do
  current_time=$(date +%s)
  elapsed=$((current_time - start_time))
  
  # 检查进程是否存活
  if ! ps -p $VLLM_PID > /dev/null; then
    echo "错误：vLLM服务进程已意外终止！"
    echo "日志内容："
    tail -n 50 $LOG_DIR/vllm_$(date +%Y%m%d).log
    exit 1
  fi
  
  # 尝试连接API服务 - 使用认证头
  if curl -s -H "Authorization: Bearer $API_KEY" http://localhost:$PORT_VLLM/v1/models > /dev/null; then
    echo "vLLM服务已成功启动！可通过 http://localhost:$PORT_VLLM 访问"
    break
  fi
  
  # 检查加载进度
  check_loading_progress
  
  # 检查是否完成启动
  if check_startup_complete; then
    echo "检测到服务已完成启动，尝试连接..."
    # 测试API连接，显示详细信息
    echo "测试API连接..."
    curl -v -H "Authorization: Bearer $API_KEY" http://localhost:$PORT_VLLM/v1/models
    sleep 5
    continue
  fi
  
  # 检查是否超时
  if [ $elapsed -ge $max_wait_time ]; then
    echo "等待服务启动超时（已等待${elapsed}秒）"
    echo "日志最后50行:"
    tail -n 50 $LOG_DIR/vllm_$(date +%Y%m%d).log
    exit 1
  fi
  
  # 显示已等待时间
  echo "已等待 ${elapsed} 秒，继续等待..."
  sleep 10
done

# 2. 启动FastAPI Web界面
echo "正在启动FastAPI Web界面，连接到vLLM服务..."

# 确保API_KEY参数正确传递
python SFT_web/Fastapi_vllm.py \
    --adapter_path $ADAPTER_PATH \
    --port $PORT_WEB \
    --vllm_url "http://localhost:$PORT_VLLM" \
    --vllm_api_key "$API_KEY" \
    > $LOG_DIR/web_$(date +%Y%m%d).log 2>&1 &

WEB_PID=$!
echo $WEB_PID > $LOG_DIR/web.pid
echo "FastAPI Web界面已启动，进程ID: $WEB_PID"
echo "日志保存在: $LOG_DIR/web_$(date +%Y%m%d).log"
echo "请访问 http://localhost:$PORT_WEB 使用Web界面"

# 可以增加对FastAPI进程的检测
if ! ps -p $WEB_PID > /dev/null; then
    echo "警告：FastAPI Web界面进程已终止！"
    echo "查看日志获取更多信息: $LOG_DIR/web_$(date +%Y%m%d).log"
fi

echo "同时按Ctrl+C退出..."
# 监听Ctrl+C信号，优雅关闭服务
trap "echo '正在关闭服务...'; kill $WEB_PID 2>/dev/null; kill $VLLM_PID 2>/dev/null; echo '服务已关闭'; exit" INT

# 保持脚本运行并显示日志
tail -f $LOG_DIR/vllm_$(date +%Y%m%d).log

# 检查端口是否正在监听
netstat -tulpn | grep 7860